import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
import matplotlib.pyplot as plt
import numpy as np
import tkinter as tk
from PIL import Image, ImageDraw
import os
import time

# 配置中文显示
import matplotlib
matplotlib.rcParams['font.sans-serif'] = ['SimHei']
matplotlib.rcParams['axes.unicode_minus'] = False

# 检查是否有可用的GPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"使用设备: {device}")

# 定义LeNet-5网络模型
class LeNet5(nn.Module):
    def __init__(self):
        super(LeNet5, self).__init__()
        # 第一个卷积层 (1x32x32 -> 6x28x28)
        self.conv1 = nn.Conv2d(1, 6, kernel_size=5)
        # 第一个池化层 (6x28x28 -> 6x14x14)
        self.pool1 = nn.AvgPool2d(kernel_size=2, stride=2)
        # 第二个卷积层 (6x14x14 -> 16x10x10)
        self.conv2 = nn.Conv2d(6, 16, kernel_size=5)
        # 第二个池化层 (16x10x10 -> 16x5x5)
        self.pool2 = nn.AvgPool2d(kernel_size=2, stride=2)
        # 第三个卷积层 (16x5x5 -> 120x1x1)
        self.conv3 = nn.Conv2d(16, 120, kernel_size=5)
        # 全连接层1 (120 -> 84)
        self.fc1 = nn.Linear(120, 84)
        # 全连接层2 (84 -> 10)
        self.fc2 = nn.Linear(84, 10)
    
    def forward(self, x):
        # 确保输入是32x32
        if x.size(2) != 32:
            # 进行上采样到32x32
            x = F.interpolate(x, size=(32, 32), mode='bilinear', align_corners=True)
        
        # 第一个卷积+激活+池化
        x = F.relu(self.conv1(x))
        x = self.pool1(x)
        
        # 第二个卷积+激活+池化
        x = F.relu(self.conv2(x))
        x = self.pool2(x)
        
        # 第三个卷积+激活
        x = F.relu(self.conv3(x))
        # 展平特征图，只保留通道维度 (batch_size, 120, 1, 1) -> (batch_size, 120)
        x = x.view(-1, 120)
        
        # 全连接层1
        x = F.relu(self.fc1(x))
        # 全连接层2（输出层）
        x = self.fc2(x)
        
        return x

# 准备数据加载
def load_data(batch_size=64):
    # 定义数据转换：将图像转换为张量并进行归一化
    transform = transforms.Compose([
        transforms.Resize((32, 32)),  # LeNet-5期望32x32输入
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,))  # MNIST数据集的均值和标准差
    ])
    
    # 加载训练集和测试集
    train_dataset = datasets.MNIST('./data', train=True, download=True, transform=transform)
    test_dataset = datasets.MNIST('./data', train=False, transform=transform)
    
    # 创建数据加载器
    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
    test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)
    
    return train_loader, test_loader

# 训练模型
def train_model(model, train_loader, optimizer, criterion, epoch, log_interval=100):
    model.train()  # 设置为训练模式
    train_loss = 0
    correct = 0
    total = 0
    
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)
        
        # 清零梯度
        optimizer.zero_grad()
        
        # 前向传播
        output = model(data)
        
        # 计算损失
        loss = criterion(output, target)
        
        # 反向传播
        loss.backward()
        
        # 更新参数
        optimizer.step()
        
        # 累计损失
        train_loss += loss.item()
        
        # 计算准确率
        _, predicted = output.max(1)
        total += target.size(0)
        correct += predicted.eq(target).sum().item()
        
        # 定期打印训练进度
        if batch_idx % log_interval == 0:
            print(f'训练 Epoch: {epoch} [{batch_idx * len(data)}/{len(train_loader.dataset)} '
                  f'({100. * batch_idx / len(train_loader):.0f}%)]\t'
                  f'损失: {loss.item():.6f}\t'
                  f'准确率: {100. * correct / total:.2f}%')
    
    # 返回平均损失和准确率
    return train_loss / len(train_loader), 100. * correct / total

# 测试模型
def test_model(model, test_loader, criterion):
    model.eval()  # 设置为评估模式
    test_loss = 0
    correct = 0
    
    with torch.no_grad():  # 不计算梯度
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)
            
            # 前向传播
            output = model(data)
            
            # 计算损失
            test_loss += criterion(output, target).item()
            
            # 获取预测结果
            _, predicted = output.max(1)
            
            # 计算正确预测数量
            correct += predicted.eq(target).sum().item()
    
    # 计算平均损失和准确率
    test_loss /= len(test_loader)
    accuracy = 100. * correct / len(test_loader.dataset)
    
    print(f'\n测试集: 平均损失: {test_loss:.4f}, 准确率: {correct}/{len(test_loader.dataset)} ({accuracy:.2f}%)\n')
    
    return test_loss, accuracy

# 可视化模型训练过程
def plot_training_results(train_losses, train_accuracies, test_losses, test_accuracies):
    plt.figure(figsize=(12, 5))
    
    plt.subplot(1, 2, 1)
    plt.plot(train_losses, label='训练损失')
    plt.plot(test_losses, label='测试损失')
    plt.xlabel('Epoch')
    plt.ylabel('损失')
    plt.legend()
    plt.title('训练和测试损失')
    
    plt.subplot(1, 2, 2)
    plt.plot(train_accuracies, label='训练准确率')
    plt.plot(test_accuracies, label='测试准确率')
    plt.xlabel('Epoch')
    plt.ylabel('准确率 (%)')
    plt.legend()
    plt.title('训练和测试准确率')
    
    plt.tight_layout()
    plt.savefig('lenet5_training_results.png')
    plt.show()

# 可视化卷积层过滤器
def visualize_filters(model):
    # 获取第一个卷积层的权重
    weights = model.conv1.weight.data.cpu().numpy()
    
    plt.figure(figsize=(10, 2))
    plt.title("第一个卷积层的过滤器")
    
    for i in range(6):  # LeNet-5的第一个卷积层有6个过滤器
        plt.subplot(1, 6, i+1)
        plt.imshow(weights[i, 0], cmap='gray')
        plt.axis('off')
    
    plt.savefig('lenet5_filters.png')
    plt.show()

# 可视化特征图
def visualize_feature_maps(model, data_loader):
    # 获取一个批次的数据
    dataiter = iter(data_loader)
    images, labels = next(dataiter)
    
    # 选择一个图像进行可视化
    img = images[0:1].to(device)
    
    # 获取第一个卷积层的输出
    model.eval()
    with torch.no_grad():
        # 确保输入是32x32
        if img.size(2) != 32:
            img = F.interpolate(img, size=(32, 32), mode='bilinear', align_corners=True)
        
        # 获取第一个卷积层的输出
        conv1_output = F.relu(model.conv1(img))
        
        # 获取第一个池化层的输出
        pool1_output = model.pool1(conv1_output)
        
        # 获取第二个卷积层的输出
        conv2_output = F.relu(model.conv2(pool1_output))
    
    # 可视化原始图像和特征图
    plt.figure(figsize=(15, 10))
    
    # 显示原始图像
    plt.subplot(1, 3, 1)
    plt.imshow(images[0].squeeze().cpu().numpy(), cmap='gray')
    plt.title(f"原始图像: {labels[0].item()}")
    plt.axis('off')
    
    # 显示第一个卷积层的特征图
    plt.subplot(1, 3, 2)
    # 将特征图重新排列为一个网格
    feature_maps = conv1_output.squeeze().cpu().numpy()
    grid_size = int(np.ceil(np.sqrt(feature_maps.shape[0])))
    fig_conv1 = plt.gcf()
    for i in range(feature_maps.shape[0]):
        plt.subplot(grid_size, grid_size, i+1)
        plt.imshow(feature_maps[i], cmap='viridis')
        plt.axis('off')
    plt.tight_layout()
    plt.title("第一个卷积层特征图")
    
    # 显示第二个卷积层的特征图
    plt.subplot(1, 3, 3)
    # 将特征图重新排列为一个网格
    feature_maps = conv2_output.squeeze().cpu().numpy()
    grid_size = int(np.ceil(np.sqrt(feature_maps.shape[0])))
    fig_conv2 = plt.gcf()
    for i in range(min(16, feature_maps.shape[0])):  # 只显示前16个特征图
        plt.subplot(4, 4, i+1)
        plt.imshow(feature_maps[i], cmap='viridis')
        plt.axis('off')
    plt.tight_layout()
    plt.title("第二个卷积层特征图")
    
    plt.savefig('lenet5_feature_maps.png')
    plt.show()

# 绘图界面类
class DrawingApp:
    def __init__(self, model):
        self.model = model
        self.model.eval()  # 设置为评估模式
        
        self.root = tk.Tk()
        self.root.title("手写数字识别 - LeNet-5")
        
        # 设置窗口大小
        self.root.geometry("600x450")
        
        # 创建画布框架
        self.canvas_frame = tk.Frame(self.root, bd=2, relief=tk.SUNKEN)
        self.canvas_frame.grid(row=0, column=0, padx=10, pady=10)
        
        # 创建画布
        self.canvas = tk.Canvas(self.canvas_frame, width=280, height=280, bg="black")
        self.canvas.pack()
        
        # 绑定鼠标事件
        self.canvas.bind("<B1-Motion>", self.paint)
        
        # 创建控制按钮框架
        self.control_frame = tk.Frame(self.root)
        self.control_frame.grid(row=0, column=1, padx=10, pady=10)
        
        # 识别结果标签
        self.result_label = tk.Label(self.control_frame, text="绘制一个数字", font=("Helvetica", 16))
        self.result_label.pack(pady=10)
        
        # 预测结果
        self.prediction_text = tk.StringVar()
        self.prediction_text.set("预测结果: ")
        self.prediction_label = tk.Label(self.control_frame, textvariable=self.prediction_text, font=("Helvetica", 14))
        self.prediction_label.pack(pady=5)
        
        # 概率条形图框架
        self.probs_frame = tk.Frame(self.control_frame)
        self.probs_frame.pack(pady=10, fill=tk.X)
        
        # 概率标签
        self.prob_labels = []
        for i in range(10):
            frame = tk.Frame(self.probs_frame)
            frame.pack(fill=tk.X, pady=2)
            
            # 数字标签
            digit_label = tk.Label(frame, text=f"{i}:", width=2)
            digit_label.pack(side=tk.LEFT, padx=5)
            
            # 概率条
            prob_bar = tk.Canvas(frame, width=150, height=15, bg="white", bd=1, relief=tk.SUNKEN)
            prob_bar.pack(side=tk.LEFT)
            
            # 概率值标签
            prob_value = tk.Label(frame, text="0%", width=5)
            prob_value.pack(side=tk.LEFT, padx=5)
            
            self.prob_labels.append((prob_bar, prob_value))
        
        # 清除按钮
        self.clear_button = tk.Button(self.control_frame, text="清除", command=self.clear_canvas)
        self.clear_button.pack(pady=10)
        
        # 识别按钮
        self.recognize_button = tk.Button(self.control_frame, text="识别", command=self.recognize)
        self.recognize_button.pack(pady=5)
        
        # 退出按钮
        self.quit_button = tk.Button(self.control_frame, text="退出", command=self.root.quit)
        self.quit_button.pack(pady=5)
        
        # 画笔宽度
        self.line_width = 15
        
        # 上一个绘制点
        self.x = None
        self.y = None
        
        # 保存用于识别的图像
        self.image = Image.new("L", (280, 280), "black")
        self.draw = ImageDraw.Draw(self.image)
    
    def paint(self, event):
        # 绘制线条
        if self.x is not None and self.y is not None:
            self.canvas.create_line(self.x, self.y, event.x, event.y,
                                   fill="white", width=self.line_width, capstyle=tk.ROUND, smooth=tk.TRUE)
            self.draw.line([self.x, self.y, event.x, event.y], fill="white", width=self.line_width)
        
        self.x = event.x
        self.y = event.y
    
    def clear_canvas(self):
        # 清除画布
        self.canvas.delete("all")
        self.image = Image.new("L", (280, 280), "black")
        self.draw = ImageDraw.Draw(self.image)
        self.x = None
        self.y = None
        
        # 重置预测结果
        self.prediction_text.set("预测结果: ")
        
        # 重置概率条
        for prob_bar, prob_value in self.prob_labels:
            prob_bar.delete("all")
            prob_value.config(text="0%")
    
    def recognize(self):
        # 将图像调整为模型输入大小
        img = self.image.resize((28, 28))
        
        # 转换为模型所需的格式
        img_tensor = transforms.ToTensor()(img)
        img_tensor = transforms.Normalize((0.1307,), (0.3081,))(img_tensor)
        img_tensor = img_tensor.unsqueeze(0).to(device)
        
        # 进行预测
        with torch.no_grad():
            output = self.model(img_tensor)
            probabilities = F.softmax(output, dim=1)[0]
            predicted_class = torch.argmax(output, dim=1).item()
        
        # 显示预测结果
        self.prediction_text.set(f"预测结果: {predicted_class}")
        
        # 更新概率条
        for i, (prob_bar, prob_value) in enumerate(self.prob_labels):
            prob = probabilities[i].item()
            width = int(prob * 150)
            
            # 清除旧条形
            prob_bar.delete("all")
            
            # 绘制新条形
            if i == predicted_class:
                color = "green"
            else:
                color = "blue"
            
            prob_bar.create_rectangle(0, 0, width, 15, fill=color, outline="")
            prob_value.config(text=f"{prob:.1%}")
    
    def run(self):
        self.root.mainloop()

def main():
    # 设置随机种子以确保结果可复现
    torch.manual_seed(42)
    
    # 超参数
    batch_size = 64
    learning_rate = 0.001
    epochs = 10
    
    # 加载数据
    print("正在加载MNIST数据集...")
    train_loader, test_loader = load_data(batch_size)
    print(f"数据加载完成! 训练集大小: {len(train_loader.dataset)}, 测试集大小: {len(test_loader.dataset)}")
    
    # 创建模型
    model = LeNet5().to(device)
    print(model)
    
    # 定义损失函数和优化器
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=learning_rate)
    
    # 检查模型文件是否存在
    model_path = 'lenet5_model.pth'
    
    if os.path.exists(model_path):
        # 加载已有模型
        print(f"找到已训练的模型 {model_path}，正在加载...")
        model.load_state_dict(torch.load(model_path, map_location=device))
        
        # 评估模型
        print("评估已加载的模型...")
        test_loss, test_accuracy = test_model(model, test_loader, criterion)
        
        # 询问用户是否要重新训练
        train_model_input = input("是否要重新训练模型? (y/n): ")
        if train_model_input.lower() != 'y':
            # 可视化滤波器和特征图
            visualize_filters(model)
            visualize_feature_maps(model, test_loader)
            
            # 启动绘图界面
            print("启动手写数字识别界面...")
            app = DrawingApp(model)
            app.run()
            return
    
    # 训练模型
    print("开始训练模型...")
    start_time = time.time()
    
    train_losses = []
    train_accuracies = []
    test_losses = []
    test_accuracies = []
    
    for epoch in range(1, epochs + 1):
        epoch_start_time = time.time()
        
        # 训练一个epoch
        train_loss, train_accuracy = train_model(model, train_loader, optimizer, criterion, epoch)
        train_losses.append(train_loss)
        train_accuracies.append(train_accuracy)
        
        # 测试模型
        test_loss, test_accuracy = test_model(model, test_loader, criterion)
        test_losses.append(test_loss)
        test_accuracies.append(test_accuracy)
        
        epoch_time = time.time() - epoch_start_time
        print(f"Epoch {epoch} 完成，用时 {epoch_time:.2f} 秒")
    
    total_time = time.time() - start_time
    print(f"模型训练完成! 总用时: {total_time:.2f} 秒")
    
    # 保存模型
    torch.save(model.state_dict(), model_path)
    print(f"模型已保存到 {model_path}")
    
    # 可视化训练结果
    plot_training_results(train_losses, train_accuracies, test_losses, test_accuracies)
    
    # 可视化滤波器和特征图
    visualize_filters(model)
    visualize_feature_maps(model, test_loader)
    
    # 启动绘图界面
    print("启动手写数字识别界面...")
    app = DrawingApp(model)
    app.run()

if __name__ == "__main__":
    main()